﻿using System;
using System.Collections.Generic;
using System.Diagnostics;

namespace g4
{
    /// <summary>
    /// 基于迭代最近点（ICP）算法的网格对齐类
    /// 功能：通过迭代计算旋转和平移变换，将源点集（Source）与目标表面（TargetSurface）对齐
    /// 核心原理：不断寻找源点在目标表面上的最近点，求解最优变换矩阵，最小化源点与对应最近点的距离误差
    /// </summary>
    public class MeshICP
    {
        /// <summary>
        /// 源点集（待对齐的点集）
        /// </summary>
        public IPointSet Source;

        /// <summary>
        /// 目标表面的AABB树（用于高效查找最近点）
        /// </summary>
        public DMeshAABBTree3 TargetSurface;

        /// <summary>
        /// 计算得到的平移变换（源点集需要施加的平移量）
        /// </summary>
        public Vector3d Translation;

        /// <summary>
        /// 计算得到的旋转变换（源点集需要施加的旋转量）
        /// </summary>
        public Quaterniond Rotation;

        /// <summary>
        ///  verbose输出回调函数（用于打印迭代过程信息）
        /// </summary>
        public Action<string> VerboseF = null;

        /// <summary>
        /// 最大迭代次数（防止无限迭代）
        /// </summary>
        public int MaxIterations = 50;

        /// <summary>
        /// 是否使用法向量辅助对齐（通过法向量方向一致性调整权重）
        /// </summary>
        public bool UseNormals = false;

        /// <summary>
        /// 最大允许距离（超过此距离的点对将被忽略）
        /// </summary>
        public double MaxAllowableDistance = double.MaxValue;

        /// <summary>
        /// 收敛容差（连续迭代的误差变化小于此值时认为收敛）
        /// </summary>
        public double ConvergeTolerance = 0.00001;

        /// <summary>
        /// 标记是否已收敛
        /// </summary>
        public bool Converged = false;

        // 内部状态变量
        private bool is_initialized = false;  // 是否已初始化内部数据结构
        private int[] MapV;                  // 源点集顶点ID到内部索引的映射
        private Vector3d[] From;             // 存储源点经当前变换后的坐标
        private Vector3d[] To;               // 存储目标表面上对应的最近点坐标
        private double[] Weights;            // 每个点对的权重（影响变换求解）
        private double LastError;            // 上一次迭代的误差


        /// <summary>
        /// 构造函数：初始化ICP对齐器
        /// </summary>
        /// <param name="source">待对齐的源点集</param>
        /// <param name="target">目标表面的AABB树（用于高效最近点查询）</param>
        public MeshICP(IPointSet source, DMeshAABBTree3 target)
        {
            Source = source;
            TargetSurface = target;

            // 初始变换：无平移、无旋转
            Translation = Vector3d.Zero;
            Rotation = Quaterniond.Identity;
        }


        /// <summary>
        /// 执行ICP迭代求解（最多MaxIterations次，或直到收敛）
        /// </summary>
        /// <param name="bUpdate">是否在已有变换基础上继续优化（true则基于当前Translation/Rotation迭代，false则重置变换）</param>
        public void Solve(bool bUpdate = false)
        {
            // 若不基于已有结果，则重置初始化状态
            if (bUpdate == false)
                is_initialized = false;
            // 首次执行或重置后，初始化内部数据结构
            if (is_initialized == false)
            {
                initialize();
                is_initialized = true;
            }

            // 应用当前变换到源点，更新From数组
            update_from();
            // 为每个源点找到目标表面上的最近点，更新To数组
            update_to();

            // 计算初始误差
            LastError = measure_error();

            // 收敛判断计数器：连续nMaxTolPassed次迭代误差变化小于容差则认为收敛
            int nTolPasssed = 0;
            int nMaxTolPassed = 5;

            // 迭代主循环
            for (int i = 0; i < MaxIterations && nTolPasssed < nMaxTolPassed; ++i)
            {
                // 打印verbose信息
                if (VerboseF != null)
                    VerboseF(string.Format("[ICP] 迭代 {0} : 误差 {1}", i, LastError));

                // 求解新的变换（旋转+平移）
                update_transformation();
                // 应用新变换到源点
                update_from();
                // 重新查找最近点
                update_to();

                // 计算当前误差
                double err = measure_error();
                // 判断误差变化是否小于容差
                if (Math.Abs(LastError - err) < ConvergeTolerance)
                {
                    nTolPasssed++;
                }
                else
                {
                    LastError = err;
                    nTolPasssed = 0;  // 误差变化较大，重置计数器
                }
            }

            // 标记是否收敛
            Converged = (nTolPasssed >= nMaxTolPassed);
        }


        /// <summary>
        /// 获取最后一次测量的误差（当前为加权平均距离）
        /// </summary>
        public double Error
        {
            get { return LastError; }
        }


        /// <summary>
        /// 将计算得到的变换应用到可变形网格的顶点（更新顶点位置和法向量）
        /// </summary>
        /// <param name="target">待更新的可变形网格</param>
        public void UpdateVertices(IDeformableMesh target)
        {
            bool hasNormals = target.HasVertexNormals;

            // 应用当前变换到源点，更新From数组
            update_from();
            // 遍历网格顶点，更新位置（和法向量）
            foreach (int vid in target.VertexIndices())
            {
                int i = MapV[vid];  // 映射到内部索引
                target.SetVertex(vid, From[i]);  // 更新顶点位置

                // 若网格有法向量，同步更新法向量（应用旋转）
                if (hasNormals)
                {
                    target.SetVertexNormal(vid,
                        (Vector3f)(Rotation * target.GetVertexNormal(vid)));
                }
            }
        }


        /// <summary>
        /// 初始化内部数据结构（源点映射、坐标数组、权重数组等）
        /// </summary>
        private void initialize()
        {
            // 初始化数组（长度为源点集顶点数量）
            From = new Vector3d[Source.VertexCount];
            To = new Vector3d[Source.VertexCount];
            Weights = new double[Source.VertexCount];

            // 初始化顶点ID映射（源点集顶点ID -> 内部数组索引）
            MapV = new int[Source.MaxVertexID];

            int i = 0;
            foreach (int vid in Source.VertexIndices())
            {
                MapV[vid] = i;          // 记录映射关系
                Weights[i] = 1.0f;      // 初始权重为1
                From[i++] = Source.GetVertex(vid);  // 初始源点坐标
            }
        }


        /// <summary>
        /// 应用当前旋转和平移变换到源点，更新From数组
        /// </summary>
        private void update_from()
        {
            int i = 0;
            foreach (int vid in Source.VertexIndices())
            {
                Weights[i] = 1.0f;  // 重置权重
                Vector3d v = Source.GetVertex(vid);  // 原始源点坐标

                // 应用变换：旋转后平移（From[i] = Rotation * v + Translation）
                From[i++] = (Rotation * v) + Translation;
            }
        }

        /// <summary>
        /// 为每个变换后的源点（From[i]）查找目标表面上的最近点，更新To数组和权重
        /// </summary>
        private void update_to()
        {
            double max_dist_sq = MaxAllowableDistance * MaxAllowableDistance;  // 最大允许距离的平方（避免开方运算）
            bool useNormals = (UseNormals && Source.HasVertexNormals);  // 是否使用法向量权重

            // 并行处理每个点（提高效率）
            Interval1i range = Interval1i.Range(From.Length);
            gParallel.ForEach(range, (vi) => {
                // 查找目标表面上距离From[vi]最近的三角形
                int tid = TargetSurface.FindNearestTriangle(From[vi], MaxAllowableDistance);
                if (tid == DMesh3.InvalidID)
                {  // 未找到有效三角形
                    Weights[vi] = 0;
                    return;
                }

                // 计算点到三角形的距离和最近点
                DistPoint3Triangle3 dist = MeshQueries.TriangleDistance(TargetSurface.Mesh, tid, From[vi]);
                if (dist.DistanceSquared > max_dist_sq)
                {  // 距离超过最大允许值
                    Weights[vi] = 0;
                    return;
                }

                // 更新最近点坐标和基础权重
                To[vi] = dist.TriangleClosest;
                Weights[vi] = 1.0f;

                // 若使用法向量，根据法向量方向一致性调整权重
                if (useNormals)
                {
                    // 源点法向量经旋转后与目标三角形法向量的点积
                    Vector3d sourceNormal = Rotation * Source.GetVertexNormal(vi);  // 源点法向量（经旋转）
                    Vector3d targetNormal = TargetSurface.Mesh.GetTriNormal(tid);   // 目标三角形法向量
                    double dot = sourceNormal.Dot(targetNormal);

                    Debug.Assert(MathUtil.IsFinite(dot));  // 确保点积有效
                    if (dot < 0)
                    {  // 法向量方向相反，权重设为0（排除该点对）
                        Weights[vi] = 0;
                    }
                    else
                    {  // 方向一致，权重增加（强化该点对的影响）
                        Weights[vi] += Math.Sqrt(dot);
                    }
                }
            });
        }


        /// <summary>
        /// 计算当前源点（From）与目标点（To）的误差（加权平均距离）
        /// </summary>
        /// <returns>误差值</returns>
        private double measure_error()
        {
            double sum = 0;    // 加权距离和
            double wsum = 0;   // 权重和
            for (int i = 0; i < From.Length; ++i)
            {
                sum += Weights[i] * From[i].Distance(To[i]);  // 权重×距离
                wsum += Weights[i];                           // 累加权重
            }
            return sum / wsum;  // 加权平均距离
        }


        /// <summary>
        /// 求解最优变换（旋转+平移），最小化From到To的距离误差
        /// 基于奇异值分解（SVD）计算最优旋转矩阵，再推导平移向量
        /// 注意：调用后From和To数组会失效，需重新调用update_from和update_to更新
        /// </summary>
        private void update_transformation()
        {
            int N = From.Length;

            // 归一化权重（权重和为1）
            double wSum = 0;
            for (int i = 0; i < N; ++i)
            {
                wSum += Weights[i];
            }
            double wSumInv = 1.0 / wSum;  // 权重和的倒数

            // 计算源点集（From）和目标点集（To）的加权均值
            Vector3d meanFrom = Vector3d.Zero;
            Vector3d meanTo = Vector3d.Zero;
            for (int i = 0; i < N; ++i)
            {
                double weight = Weights[i] * wSumInv;  // 归一化权重
                meanFrom += weight * From[i];
                meanTo += weight * To[i];
            }

            // 去均值（将点集中心移到原点，消除平移影响）
            for (int i = 0; i < N; ++i)
            {
                From[i] -= meanFrom;
                To[i] -= meanTo;
            }

            // 构造协方差矩阵M（3x3）：M = From * To^T（加权）
            double[] M = new double[9];  // 按行存储：[M00, M01, M02, M10, M11, M12, M20, M21, M22]
            for (int k = 0; k < 3; ++k)
            {  // k：行索引（对应From的x/y/z分量）
                int rowStart = 3 * k;  // 当前行在数组中的起始索引
                for (int i = 0; i < N; ++i)
                {
                    double lhs = Weights[i] * wSumInv * From[i][k];  // 加权的From分量
                    M[rowStart + 0] += lhs * To[i].x;  // 列0（To的x分量）
                    M[rowStart + 1] += lhs * To[i].y;  // 列1（To的y分量）
                    M[rowStart + 2] += lhs * To[i].z;  // 列2（To的z分量）
                }
            }

            // 对协方差矩阵M进行奇异值分解（SVD）：M = U * S * V^T
            SingularValueDecomposition svd = new SingularValueDecomposition(3, 3, 100);
            uint ok = svd.Solve(M, -1);  // 求解SVD，按奇异值降序排列
            Debug.Assert(ok < 9999999);  // 确保SVD求解成功

            // 获取左奇异矩阵U和右奇异矩阵V
            double[] U = new double[9], V = new double[9], temp = new double[9];
            svd.GetU(U);
            svd.GetV(V);

            // 计算最优旋转矩阵（基于SVD结果）
            double[] rotUpdate = new double[9];  // 旋转矩阵更新量

            // 确保旋转矩阵行列式为1（避免镜像变换）
            double detU = MatrixUtil.Determinant3x3(U);  // U的行列式
            double detV = MatrixUtil.Determinant3x3(V);  // V的行列式
            if (detU * detV < 0)
            {  // 行列式乘积为负，存在镜像，修正V
                double[] S = MatrixUtil.MakeDiagonal3x3(1, 1, -1);  // 对角矩阵（最后一个元素为-1）
                MatrixUtil.Multiply3x3(V, S, temp);  // V = V * S（修正V）
                MatrixUtil.Transpose3x3(U);          // U^T
                MatrixUtil.Multiply3x3(temp, U, rotUpdate);  // 旋转矩阵 = V*S * U^T
            }
            else
            {  // 行列式乘积为正，直接计算
                MatrixUtil.Transpose3x3(U);          // U^T
                MatrixUtil.Multiply3x3(V, U, rotUpdate);  // 旋转矩阵 = V * U^T
            }

            // 将旋转矩阵转换为四元数
            Matrix3d rotUpdateMat = new Matrix3d(rotUpdate);
            Quaterniond rotUpdateQuat = new Quaterniond(rotUpdateMat);

            // 计算平移更新量：平移 = 目标均值 - 旋转×源均值
            Vector3d transUpdate = meanTo - rotUpdateQuat * meanFrom;
            Translation += transUpdate;  // 更新总平移

            // 更新总旋转（累积旋转）
            Rotation = rotUpdateQuat * Rotation;
        }
    }
}